from typing import List, Dict, Optional
from datetime import datetime
from pydantic import BaseModel
import json

class ChatHistory(BaseModel):
    session_id: str
    messages: List[Dict[str, str]]
    created_at: datetime
    updated_at: datetime

class MemoryService:
    def __init__(self):
        self.sessions: Dict[str, ChatHistory] = {}
    
    def create_session(self, session_id: str) -> ChatHistory:
        """创建新的会话"""
        now = datetime.now()
        history = ChatHistory(
            session_id=session_id,
            messages=[],
            created_at=now,
            updated_at=now
        )
        self.sessions[session_id] = history
        return history
    
    def add_message(self, session_id: str, role: str, content: str) -> None:
        """添加消息到会话历史"""
        if session_id not in self.sessions:
            self.create_session(session_id)
        
        history = self.sessions[session_id]
        history.messages.append({"role": role, "content": content})
        history.updated_at = datetime.now()
    
    def get_messages(self, session_id: str, limit: int = 20) -> List[Dict[str, str]]:
        """获取会话历史消息，默认保留20轮对话"""
        if session_id not in self.sessions:
            return []
        return self.sessions[session_id].messages[-limit:]
    
    def get_last_user_message(self, session_id: str) -> Optional[str]:
        """获取上一条用户消息"""
        if session_id not in self.sessions:
            return None
        
        # 从后向前查找最近的一条用户消息
        for msg in reversed(self.sessions[session_id].messages):
            if msg["role"] == "user":
                return msg["content"]
        
        return None
    
    def clear_session(self, session_id: str) -> None:
        """清除会话历史"""
        if session_id in self.sessions:
            del self.sessions[session_id]
    
    def clear_history(self, session_id: str) -> None:
        """清除历史记录（clear_session的别名）"""
        self.clear_session(session_id)
    
    def get_history(self, session_id: str, limit: int = 20) -> List[Dict[str, str]]:
        """获取历史记录（get_messages的别名）"""
        return self.get_messages(session_id, limit)
    
    def get_session_count(self) -> int:
        """获取当前会话数量"""
        return len(self.sessions)
    
    def get_all_sessions(self) -> List[str]:
        """获取所有会话ID"""
        return list(self.sessions.keys())
    
    def session_exists(self, session_id: str) -> bool:
        """检查会话是否存在"""
        return session_id in self.sessions 